//! 移位指令编码模块
//! 
//! 实现 8086 移位指令：SHL, SHR, SAR

use crate::ast::Operand;
use crate::error::{AssemblerError, Result};
use super::registers::RegisterEncoding;

/// 移位指令编码器
pub struct ShiftEncoder;

impl ShiftEncoder {
    /// 编码移位指令
    pub fn encode(&self, instruction: &str, operands: &[Operand]) -> Result<Vec<u8>> {
        match instruction {
            "左移" => self.encode_shift(operands, 0b100, "左移"),
            "右移" => self.encode_shift(operands, 0b101, "右移"),
            "算术右移" => self.encode_shift(operands, 0b111, "算术右移"),
            _ => Err(AssemblerError::CodeGenError {
                message: format!("未知的移位指令: {}", instruction),
            }),
        }
    }

    /// 移位指令的通用编码函数
    fn encode_shift(
        &self,
        operands: &[Operand],
        shift_type: u8, // ModR/M reg字段：SHL=100, SHR=101, SAR=111
        instruction_name: &str,
    ) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: format!("{} 指令需要 2 个操作数", instruction_name),
                line: 0,
            });
        }

        let reg = match &operands[0] {
            Operand::Register(reg) => reg,
            _ => {
                return Err(AssemblerError::OperandError {
                    message: format!("{} 第一个操作数必须是寄存器", instruction_name),
                    line: 0,
                });
            }
        };

        let reg_code = RegisterEncoding::get_16bit_register_code(reg)
            .ok_or_else(|| AssemblerError::OperandError {
                message: format!("无效的16位寄存器: {}", reg),
                line: 0,
            })?;

        match &operands[1] {
            Operand::Immediate(count) => {
                if *count == 1 {
                    // 移位1位：D1 /r (16位)
                    let opcode = 0xD1;
                    let modrm = (0b11 << 6) | (shift_type << 3) | reg_code;
                    Ok(vec![opcode, modrm])
                } else if *count > 0 && *count <= 255 {
                    // 移位imm8位：C1 /r imm8 (16位)
                    let opcode = 0xC1;
                    let modrm = (0b11 << 6) | (shift_type << 3) | reg_code;
                    Ok(vec![opcode, modrm, *count as u8])
                } else {
                    Err(AssemblerError::OperandError {
                        message: format!(
                            "{} 移位次数必须在 1 到 255 之间，得到: {}",
                            instruction_name, count
                        ),
                        line: 0,
                    })
                }
            }
            Operand::Register(count_reg) => {
                // 检查是否是CL寄存器（计数寄存器）
                if count_reg == "计数" {
                    // 移位CL位：D3 /r (16位)
                    let opcode = 0xD3;
                    let modrm = (0b11 << 6) | (shift_type << 3) | reg_code;
                    Ok(vec![opcode, modrm])
                } else {
                    Err(AssemblerError::OperandError {
                        message: format!(
                            "{} 第二个操作数必须是立即数或计数寄存器，得到: {}",
                            instruction_name, count_reg
                        ),
                        line: 0,
                    })
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: format!("{} 第二个操作数必须是立即数或计数寄存器", instruction_name),
                line: 0,
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_shl_reg_1() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(1),
        ];
        let result = encoder.encode("左移", &operands).unwrap();
        // SHL AX, 1 = D1 E0
        assert_eq!(result, vec![0xD1, 0xE0]);
    }

    #[test]
    fn test_encode_shl_reg_imm8() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(5),
        ];
        let result = encoder.encode("左移", &operands).unwrap();
        // SHL AX, 5 = C1 E0 05
        assert_eq!(result, vec![0xC1, 0xE0, 0x05]);
    }

    #[test]
    fn test_encode_shl_reg_cl() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("计数".to_string()),
        ];
        let result = encoder.encode("左移", &operands).unwrap();
        // SHL AX, CL = D3 E0
        assert_eq!(result, vec![0xD3, 0xE0]);
    }

    #[test]
    fn test_encode_shr_reg_1() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(1),
        ];
        let result = encoder.encode("右移", &operands).unwrap();
        // SHR AX, 1 = D1 E8
        assert_eq!(result, vec![0xD1, 0xE8]);
    }

    #[test]
    fn test_encode_sar_reg_1() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(1),
        ];
        let result = encoder.encode("算术右移", &operands).unwrap();
        // SAR AX, 1 = D1 F8
        assert_eq!(result, vec![0xD1, 0xF8]);
    }

    #[test]
    fn test_encode_shift_invalid_count() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(0),
        ];
        let result = encoder.encode("左移", &operands);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_shift_invalid_register() {
        let encoder = ShiftEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("左移", &operands);
        assert!(result.is_err());
    }
}

